#include "byte_stream.hh"
#include "stream_reassembler.hh"
#include "util.hh"

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <exception>
#include <iostream>
#include <stdexcept>
#include <tuple>
#include <utility>
#include <vector>

using namespace std;

static constexpr unsigned NREPS = 32;
static constexpr unsigned NSEGS = 128;
static constexpr unsigned MAX_SEG_LEN = 2048;

string read(StreamReassembler &reassembler) {
  return reassembler.stream_out().read(reassembler.stream_out().buffer_size());
}

int main() {
  try {
    auto rd = get_random_generator();

    // buffer a bunch of bytes, make sure we can empty and re-fill before calling close()
    for (unsigned rep_no = 0; rep_no < NREPS; ++rep_no) {
      StreamReassembler buf{MAX_SEG_LEN * NSEGS};

      vector<tuple<size_t, size_t>> seq_size;
      size_t offset = 0;
      for (unsigned i = 0; i < NSEGS; ++i) {
        const size_t size = 1 + (rd() % (MAX_SEG_LEN - 1));
        seq_size.emplace_back(offset, size);
        offset += size;
      }
      shuffle(seq_size.begin(), seq_size.end(), rd);

      string d(offset, 0);
      generate(d.begin(), d.end(), [&] { return rd(); });

      for (auto [off, sz] : seq_size) {
        string dd(d.cbegin() + off, d.cbegin() + off + sz);
        buf.push_substring(move(dd), off, off + sz == offset);
      }

      auto result = read(buf);
      if (buf.stream_out().bytes_written() != offset) {  // read bytes
        throw runtime_error("test 1 - number of bytes RX is incorrect");
      }
      if (!equal(result.cbegin(), result.cend(), d.cbegin())) {
        throw runtime_error("test 1 - content of RX bytes is incorrect");
      }
    }

    // insert EOF into a hole in the buffer
    for (unsigned rep_no = 0; rep_no < NREPS; ++rep_no) {
      StreamReassembler buf{65'000};

      const size_t size = 1024;
      string d(size, 0);
      generate(d.begin(), d.end(), [&] { return rd(); });

      buf.push_substring(d, 0, false);
      buf.push_substring(d.substr(10), size + 10, false);

      auto res1 = read(buf);
      if (buf.stream_out().bytes_written() != size) {
        throw runtime_error("test 3 - number of RX bytes is incorrect");
      }
      if (!equal(res1.cbegin(), res1.cend(), d.cbegin())) {
        throw runtime_error("test 3 - content of RX bytes is incorrect");
      }

      buf.push_substring(string(d.cbegin(), d.cbegin() + 7), size, false);
      buf.push_substring(string(d.cbegin() + 7, d.cbegin() + 8), size + 7, true);

      auto res2 = read(buf);
      if (buf.stream_out().bytes_written() != size + 8) {  // rx bytes
        throw runtime_error("test 3 - number of RX bytes is incorrect after 2nd read");
      }
      if (!equal(res2.cbegin(), res2.cend(), d.cbegin())) {
        throw runtime_error("test 3 - content of RX bytes is incorrect after 2nd read");
      }
    }

    // insert EOF over previously queued data, require one of two possible correct actions
    for (unsigned rep_no = 0; rep_no < NREPS; ++rep_no) {
      StreamReassembler buf{65'000};

      const size_t size = 1024;
      string d(size, 0);
      generate(d.begin(), d.end(), [&] { return rd(); });

      buf.push_substring(d, 0, false);
      buf.push_substring(d.substr(10), size + 10, false);

      auto res1 = read(buf);
      if (buf.stream_out().bytes_written() != size) {
        throw runtime_error("test 4 - number of RX bytes is incorrect");
      }
      if (!equal(res1.cbegin(), res1.cend(), d.cbegin())) {
        throw runtime_error("test 4 - content of RX bytes is incorrect");
      }

      buf.push_substring(string(d.cbegin(), d.cbegin() + 15), size, true);

      auto res2 = read(buf);
      if (buf.stream_out().bytes_written() != 2 * size && buf.stream_out().bytes_written() != size + 15) {
        throw runtime_error("test 4 - number of RX bytes is incorrect after 2nd read");
      }
      if (!equal(res2.cbegin(), res2.cend(), d.cbegin())) {
        throw runtime_error("test 4 - content of RX bytes is incorrect after 2nd read");
      }
    }
  } catch (const exception &e) {
    cerr << "Exception: " << e.what() << endl;
    return EXIT_FAILURE;
  }

  return EXIT_SUCCESS;
}
